package com.jboltai.capability.vdb.collection.field;


import com.alibaba.fastjson.JSONObject;
import com.jboltai.capability.vdb.collection.VDBMetricType;

public class VDBField {

    /**
     * 考虑到兼容性问题，向量字段名称强制为vector。因为腾讯向量数据库，强制该字段强制要求为vector
     */
    public static final String VECTOR_FIELD_NAME = "vector";

    /**
     * 强制id字段的名字
     */
    public static final String ID_FIELD_NAME = "id";

    /**
     * 强制文本字段的名字
     */
    public static final String DOC_FIELD_NAME = "doc";
    /**
     * 考虑到兼容性问题，向量字段名称强制为vector
     */
    private String name;

    private VDBFieldType vdbFieldType;

    private boolean isPrimaryKey;

    /**
     * 向量索引类型
     */
    private VectorIndexType vectorIndexType;

    /**
     * 向量索引参数
     */
    private JSONObject vectorIndexParams;

    /**
     * 指定向量索引的相似度度量算法
     */
    private VDBMetricType vdbMetricType = VDBMetricType.DEFAULT;

    /**
     * 向量维度，默认1024。维度建议为4的整数倍，字节对齐有助于提升搜索性能。维度越高，存储成本越高，检索效率越低。
     */
    private int dimension = 1024;

    /**
     * 是否支持全文检索
     */
    private boolean fullTextSearch = false;



    public VDBField() {}

    public VDBField(String name) {
        this.name = name;
    }

    public VDBField(String name, VDBFieldType vdbFieldType) {
        this.name = name;
        if (name.equals("id")) {
            this.isPrimaryKey = true;
        }
        this.vdbFieldType = vdbFieldType;
    }

    /**
     * 处于兼容性考虑，暂时不对外开放，主键字段只能使用内置的
     * @param name
     * @param vdbFieldType
     * @param isPrimaryKey
     */
    private VDBField(String name, VDBFieldType vdbFieldType, boolean isPrimaryKey) {
        this.name = name;
        this.vdbFieldType = vdbFieldType;
        this.isPrimaryKey = isPrimaryKey;
    }

    public VDBField(VectorIndexType vectorIndexType, JSONObject vectorIndexParams) {
        this.name = VECTOR_FIELD_NAME;
        this.vdbFieldType = VDBFieldType.VECTOR;
        this.vectorIndexType = vectorIndexType;
        this.vectorIndexParams = vectorIndexParams;
    }
    public VDBField(VectorIndexType vectorIndexType, JSONObject vectorIndexParams, VDBMetricType vdbMetricType) {
        this.name = VECTOR_FIELD_NAME;
        this.vdbFieldType = VDBFieldType.VECTOR;
        this.vectorIndexType = vectorIndexType;
        this.vectorIndexParams = vectorIndexParams;
        this.vdbMetricType = vdbMetricType;
    }


    /**
     * 构建一个普通字段
     * @param name
     * @param vdbFieldType
     * @return
     */
    public static VDBField ofNormal(String name, VDBFieldType vdbFieldType) {
        return new VDBField(name, vdbFieldType);
    }

//    /**
//     * 构建一个主键字段
//     * @param name
//     * @param vDbFieldType
//     * @return
//     */
//    public static VDbField ofPrimary(String name, VDbFieldType vDbFieldType) {
//        return new VDbField(name, vDbFieldType, true);
//    }

    /**
     * 默认的文本字段
     * @param fullTextSearch 是否支持全文检索，目前仅在 Milvus 中支持此特性
     * @return
     */
    public static VDBField ofDoc(boolean fullTextSearch) {
        return new VDBField(DOC_FIELD_NAME, VDBFieldType.STRING).setFullTextSearch(fullTextSearch);
    }

    /**
     * FLAT类型向量字段。<br/>
     * 因为兼容性问题，强制该字段名称为vector
     * @return
     */
    public static VDBField ofFlatVector() {
        return new VDBField(VectorIndexType.FLAT, null);
    }

    /**
     * HNSW类型向量字段<br/>因为兼容性问题，强制该字段名称为vector。<br/><br/>
     * HNSW (Hierarchical Navigable Small World Graph) 是一种基于图形的索引算法。它根据某些规则为图像构建多层导航结构。在这个结构中，上层更稀疏，节点之间的距离更远；而下层更密集，节点之间的距离更近。搜索从最上层开始，找到最接近目标的节点，在这一层中，然后进入下一层开始另一个搜索。经过多次迭代，它可以快速接近目标位置。
     * <br/><br/>
     * 为了提高性能，HNSW 限制了图的每一层节点的最大度数为M。此外，您可以使用efConstruction（在构建索引时）或ef（在搜索目标时）来指定搜索范围。
     * @param m              每个节点在检索构图中可以连接多少个邻居节点。取值范围：[4,64]。
     * @param efConstruction 搜索时，指定寻找节点邻居遍历的范围。数值越大构图效果越好，构图时间越长。取值范围：[8,512]
     * @return
     */
    public static VDBField ofHnswVector(int m, int efConstruction) {
        return new VDBField(VectorIndexType.HNSW, new JSONObject().fluentPut("M", m).fluentPut("efConstruction", efConstruction));
    }
    /**
     * HNSW类型向量字段. m值默认16，efConstruction默认200。<br/>因为兼容性问题，强制该字段名称为vector
     * @return
     */
    public static VDBField ofHnswVector() {
        return new VDBField(VectorIndexType.HNSW, new JSONObject().fluentPut("M", 16).fluentPut("efConstruction", 200));
    }


    /**
     * IVF_FLAT类型向量字段<br/>因为兼容性问题，强制该字段名称为vector。<br/><br/>
     * IVF_FLAT 将向量数据分成 nlist 个聚类单元，然后比较目标输入向量与每个聚类中心之间的距离。根据系统设置的查询聚类数目(nprobe)，只有与目标输入向量在最相似的聚类中的向量之间进行比较，才会返回相似性搜索结果——从而大大减少查询时间。
     *<br/><br/>
     * 通过调整 nprobe，可以在给定情况下找到精度和速度之间的理想平衡。来自IVF_FLAT性能测试(opens in a new tab)的结果表明，随着目标输入向量的数目(nq)和要搜索的聚类数目(nprobe)的增加，查询时间急剧增加。
     * @param nlist 指索引中的聚类中心数量。取值范围：[1,65536]。
     * @return
     */
    public static VDBField ofIvfFlatVector(int nlist) {
        return new VDBField(VectorIndexType.IVF_FLAT, new JSONObject().fluentPut("nlist", nlist));
    }

    /**
     * IVF_PQ类型向量字段。<br/>因为兼容性问题，强制该字段名称为vector。<br/><br/>
     * PQ（Product Quantization，产品量化）将原始高维向量空间等分为 m 个低维向量空间的笛卡尔积，然后对这些低维向量空间进行量化。与计算目标向量与所有单元的中心点之间的距离不同，产品量化使得可以计算目标向量与每个低维空间的聚类中心之间的距离，并大大降低了算法的时间复杂度和空间复杂度。
     * <br/><br/>
     * IVF_PQ在量化向量的积之前执行IVF索引聚类。它的索引文件甚至比IVF_SQ8更小，但也会在搜索向量时导致精度损失。
     * @param nlist 指索引中的聚类中心数量。
     * @param m 指乘积量化中原始数据被拆分的子向量的数量。取值范围：[1,向量维度]。
     * @return
     */
    public static VDBField ofIvfPqVector(int nlist, int m) {
        return new VDBField(VectorIndexType.IVF_PQ, new JSONObject().fluentPut("nlist", nlist).fluentPut("M", m));
    }

    /**
     * IVF_SQ8类型向量字段。<br/>因为兼容性问题，强制该字段名称为vector。<br/><br/>
     * IVF_FLAT不进行任何压缩，因此它生成的索引文件与原始的非索引向量数据大小大致相同。例如，如果原始的1B SIFT数据集大小为476 GB，则其IVF_FLAT索引文件大小稍大（约为470 GB）。将所有索引文件加载到内存中将消耗470 GB的存储空间。
     * <br/><br/>
     * 当磁盘、CPU或GPU内存资源受限时，IVF_SQ8比IVF_FLAT更好。该索引类型可以通过执行标量量化将每个FLOAT（4字节）转换为UINT8（1字节）。这将减少70-75%的磁盘、CPU和GPU内存消耗。对于1B SIFT数据集，IVF_SQ8索引文件只需要140 GB的存储空间。
     * @param nlist 指索引中的聚类中心数量。取值范围：[1,65536]。
     * @return
     */
    public static VDBField ofIvfSq8Vector(int nlist) {
        return new VDBField(VectorIndexType.IVF_PQ, new JSONObject().fluentPut("nlist", nlist));
    }





    public String getName() {
        return name;
    }

    public VDBField setName(String name) {
        this.name = name;
        return this;
    }

    public VDBFieldType getVdbFieldType() {
        return vdbFieldType;
    }

    public VDBField setVdbFieldType(VDBFieldType vdbFieldType) {
        this.vdbFieldType = vdbFieldType;
        return this;
    }

    public boolean isPrimaryKey() {
        return isPrimaryKey;
    }

    public VDBField setPrimaryKey(boolean primaryKey) {
        isPrimaryKey = primaryKey;
        return this;
    }

    public VectorIndexType getVectorIndexType() {
        return vectorIndexType;
    }

    public VDBField setVectorIndexType(VectorIndexType vectorIndexType) {
        this.vectorIndexType = vectorIndexType;
        return this;
    }

    public JSONObject getVectorIndexParams() {
        return vectorIndexParams;
    }

    public VDBField setVectorIndexParams(JSONObject vectorIndexParams) {
        this.vectorIndexParams = vectorIndexParams;
        return this;
    }

    public VDBMetricType getVdbMetricType() {
        return vdbMetricType;
    }

    public VDBField setVdbMetricType(VDBMetricType vdbMetricType) {
        this.vdbMetricType = vdbMetricType;
        return this;
    }

    //暂时不允许指定维度，只能由资源绑定的embedding模型决定
//    public int getDimension() {
//        return dimension;
//    }
//
//    public VDBField setDimension(int dimension) {
//        this.dimension = dimension;
//        return this;
//    }

    public boolean isFullTextSearch() {
        return fullTextSearch;
    }

    public VDBField setFullTextSearch(boolean fullTextSearch) {
        this.fullTextSearch = fullTextSearch;
        return this;
    }
}
